#include "cartographer/mapping/internal/2d/scan_matching/adaptive_matcher_angle_2d.h"

using namespace cartographer::mapping;
using namespace cartographer::mapping::scan_matching;

namespace cartographer
{
    AdaptiveFilter::AdaptiveFilter(const ScanMatchOptions &options)
    {
        inited_ = false;
        win_def_linear = options.linear_search_window();
        win_def_angular = options.angular_search_window();
        win_linear = win_def_linear;
        win_angular = win_def_angular;
        cost_weight_linear = options.translation_delta_cost_weight();
        cost_weight_angular = options.rotation_delta_cost_weight();
        LOG(INFO) << "[Adaptive] win_ang " << win_angular << " win_lin " << win_linear << " cost_weight_linear " << cost_weight_linear
                  << " cost_weight_angular " << cost_weight_angular;
    }

    void AdaptiveFilter::UpdateRotate()
    {
        if (!rotations_.empty())
            std::vector<double>().swap(rotations_);
        if (inited_)
        {
            UpdateRotateByCandidates();
        }
        else
        {
            double deltaYaw = -offsetStep(2) * angular_stride;
            for (int i = 0; i < num_scans_; ++i, deltaYaw += angular_stride)
            {
                if (std::fabs(deltaYaw) < 0.00001)
                    deltaYaw = 0.0;
                rotations_.push_back(deltaYaw);
            }
        }
        ResetLinearBounds(rotations_.size(), GetLinearStepSize());
    }

    void AdaptiveFilter::UpdateParameter(const float max_range, const double submap_resolution)
    {
        const double kSafetyMargin = 1. - 1e-3;
        // https://blog.csdn.net/weixin_44088154/article/details/133828291
        // 角度搜索步长
        angular_stride = kSafetyMargin * std::acos(1. - common::Pow2(submap_resolution) / (2. * common::Pow2(max_range)));
        // angular_stride = common::DegToRad(0.05);
        if (inited_)
        {
            Eigen::Vector3d offsetDelta = GetOffsetDelta();
            offsetStep(2) = std::ceil((offsetDelta(2) * 2.0) / angular_stride);
            if (offsetStep(2) < 1)
                offsetStep(2) = 1;
        }
        else
        {
            // 窗口内移动步数
            offsetStep(2) = std::ceil(win_angular / angular_stride);
        }
        // 扩大角度移动步数
        num_scans_ = 2 * offsetStep(2) + 1;
        UpdateRotate();
        // print message
        LOG(INFO) << "\n[AdaptiveFilter] : \n";
        LOG(INFO) << "num_scans " << num_scans_ << " map_resolution " << submap_resolution << " max_range " << max_range;
        LOG(INFO) << "angula: win " << common::RadToDeg(win_angular) << " stride " << common::RadToDeg(angular_stride);
        LOG(INFO) << "linear: win " << win_linear << " stride " << resolution_linear;
    }

    void AdaptiveFilter::ResetLinearBounds(const int num_scans, const int step_size)
    {
        if (!linear_bounds.empty())
            std::vector<LinearBounds>().swap(linear_bounds);
        for (int i = 0; i < num_scans; ++i)
            linear_bounds.emplace_back(-step_size, step_size, -step_size, step_size);
    }

    void AdaptiveFilter::UpdateRotateByCandidates()
    {
        double firstAngle = candidates[0].orientation + -offsetStep(2) * angular_stride;
        double limitMin = firstAngle;
        double limitMax = firstAngle;

        std::vector<double> orientations;
        for (size_t i = 0; i < candidates.size(); i++)
        {
            const auto &cand = candidates[i];
            // LOG(INFO) << "[KeepCand2] score " << cand.score << " yaw " << cand.orientation;
            for (int iw = -offsetStep(2); iw <= offsetStep(2); iw++)
            {
                double rotate = cand.orientation + iw * angular_stride;
                if (std::fabs(rotate) < 0.00001)
                    rotate = 0.0;
                if (rotate > limitMax)
                    limitMax = rotate;
                if (rotate < limitMin)
                    limitMin = rotate;
                orientations.push_back(rotate);
            }
        }
        limitMax += 0.0001;
        limitMin -= 0.0001;
        int groupSize = std::ceil(std::fabs(limitMax - limitMin) / angular_stride);
        if (groupSize < 1)
            return;
        std::vector<HistogramGroup> histogram;
        histogram.resize(groupSize + 1);
        for (size_t i = 0; i < orientations.size(); i++)
        {
            int gidx = std::ceil((orientations[i] - limitMin) / angular_stride);
            if (gidx > groupSize || gidx < 0)
                continue;
            histogram[gidx].count++;
            histogram[gidx].value += orientations[i];
        }
        limitMin += 0.0001;
        int sampleCnt = 0;
        for (size_t i = 0; i < histogram.size(); i++)
        {
            if (histogram[i].count < 1)
                continue;
            sampleCnt++;
            const auto &item = histogram[i];
            double yaw = item.value / double(item.count);
            // double yaw = double(i) * angular_stride + limitMin;
            rotations_.push_back(yaw);
        }
        LOG(INFO) << "[groupSize] step " << offsetStep(2) << " size " << candidates.size() << " stride " << angular_stride;
        LOG(INFO) << "[groupSize] max " << limitMax << " min " << limitMin << " first " << firstAngle;
        LOG(INFO) << "[groupSize] gpsize " << groupSize << " orisize " << orientations.size() << " sampleCnt " << sampleCnt;
    }

    void AdaptiveFilter::SampleNears()
    {
        if (!candidates.empty())
            std::vector<Candidate2D>().swap(candidates);
        CHECK_EQ(rotations_.size(), linear_bounds.size());
        for (size_t i = 0; i < rotations_.size(); ++i)
        {
            for (int xo = linear_bounds[i].min_x; xo <= linear_bounds[i].max_x; ++xo)
            {
                for (int yo = linear_bounds[i].min_y; yo <= linear_bounds[i].max_y; ++yo)
                {
                    candidates.emplace_back(candidates.size(), i, xo, yo, rotations_[i], resolution_linear);
                }
            }
        }
        weights_.resize(candidates.size(), 0.0);
        LOG(INFO) << "[SampleNears] size cand " << candidates.size() << " rots " << rotations_.size();
    }

    void AdaptiveFilter::UpdateWeightByExp()
    {
        double weight_sum = 0;
        for (size_t i = 0; i < candidates.size(); i++)
        {
            auto &cand = candidates[i];
            double weight = gauss_.Prop(cand.orientation, offsetExpect(2));
            weights_[i] = weight;
            weight_sum += weight;
        }
        NormalizeWeight(weights_, weight_sum);
    }

    int AdaptiveFilter::Resample()
    {
        double weight_max = 0;
        double weight_min = 0;
        for (size_t i = 0; i < weights_.size(); ++i)
        {
            double weight = weights_[i];
            if (i == 0 || weight_max < weight)
                weight_max = weight;
            if (i == 0 || weight_min > weight)
                weight_min = weight;
        }
        double weight_limit = (weight_max + weight_min) * 0.5;

        int groupSize = 10;
        int weight_size = 0;
        double weight_sum = 0.0;
        float groupWidth = (weight_max - weight_limit) / float(groupSize);
        std::vector<HistogramGroup> groups;
        groups.resize(groupSize + 1);

        LOG(INFO) << "weight [" << weight_max << " " << weight_limit << "] groupWidth " << groupWidth;

        for (size_t i = 0; i < weights_.size(); ++i)
        {
            double weight = weights_[i];
            if (weight < weight_limit)
                continue;
            int gpi = (weight - weight_limit) / groupWidth;
            if (gpi > groupSize || gpi < 0)
            {
                LOG(INFO) << "wrong group " << i << " index " << gpi << " weight " << weight;
                continue;
            }
            weight_sum += weight;
            weight_size++;
            auto &gpItem = groups[gpi];
            if (gpItem.id < 0)
            {
                gpItem.id = gpi;
                gpItem.count = i;
                gpItem.value = weight;
            }
            else
            {
                if (gpItem.value < weight)
                {
                    gpItem.count = i;
                    gpItem.value = weight;
                }
            }
        }

        double weight_average = weight_sum / double(weight_size);
        if (weight_slow_ < 0.0)
            weight_slow_ = weight_average;
        else
            weight_slow_ = alpha_slow_ * weight_average + (1 - alpha_slow_) * weight_slow_;

        if (weight_fast_ < 0.0)
            weight_fast_ = weight_average;
        else
            weight_fast_ = alpha_fast_ * weight_average + (1 - alpha_fast_) * weight_fast_;

        double weight_diff = 1.0 - weight_fast_ / weight_slow_;
        LOG(INFO) << "weight diff " << weight_diff << " fast " << weight_fast_ << " slow " << weight_slow_;
        if (weight_diff > 0.2)
        {
            LOG(INFO) << "filterReset";
            return -1;
        }

        std::vector<Candidate2D> keepCands;
        for (const auto &gp : groups)
        {
            if (gp.id < 0)
                continue;
            const auto &cand = candidates[gp.count];
            LOG(INFO) << "[KeepCand] score " << cand.score << " weight " << weights_[gp.count] << " yaw "
                      << common::RadToDeg(cand.orientation);
            keepCands.push_back(cand);
        }
        if (keepCands.size() < 1)
            return -2;
        LOG(INFO) << "[Resample] keep " << keepCands.size();
        // reset data
        candidates.swap(keepCands);
        weights_.clear();
        rotations_.clear();
        return 1;
    }

    int AdaptiveFilter::Success(const transform::Rigid2d &pose_init, const transform::Rigid2d &pose_opt)
    {
        offsetExpect = ComputeOffset(pose_init, pose_opt);
        Eigen::Vector3d offsetDelta = GetOffsetDelta();
        gauss_.Push(offsetExpect(2));
        LOG(INFO) << "[Feedback] org size " << candidates.size();
        LOG(INFO) << "[Feedback] offset_exp " << offsetExpect(0) << " " << offsetExpect(1) << " " << common::RadToDeg(offsetExpect(2));
        LOG(INFO) << "[Feedback] offset_est " << offsetBest(0) << " " << offsetBest(1) << " " << common::RadToDeg(offsetBest(2));
        LOG(INFO) << "[Feedback] offset_del " << offsetDelta(0) << " " << offsetDelta(1) << " " << common::RadToDeg(offsetDelta(2));
        if (candidates.empty())
            return -1;
        UpdateWeightByExp();
        // double yawDelta = common::RadToDeg(std::fabs(offsetDelta(2)));
        // if (std::fabs(offsetDelta(0)) > 0.03 || std::fabs(offsetDelta(1)) > 0.03 || yawDelta > 2.0)
        //     return -2;
        if (!Resample())
            return -3;
        if (candidates.empty())
            return -4;
        Candidate2D cand;
        cand.id = -1;
        cand.score = 0.0;
        cand.scan_index = -1;
        cand.x_index_offset = 0;
        cand.y_index_offset = 0;
        cand.x = offsetExpect(0);
        cand.y = offsetExpect(1);
        cand.orientation = offsetExpect(2);
        // candidates.push_back(cand);
        LOG(INFO) << "[Feedback] resample size " << candidates.size();
        return 1;
    }

    bool AdaptiveFilter::Check()
    {
        if (!inited_ && rotations_.size() != size_t(num_scans_))
        {
            LOG(INFO) << "[Adaptive] wrong rotations size " << rotations_.size() << " num_scans " << num_scans_;
            return false;
        }
        if (inited_ && rotations_.size() < 1)
        {
            LOG(INFO) << "[Adaptive] wrong rotations size " << rotations_.size() << " num_scans" << num_scans_;
            return false;
        }
        if (candidates.size() < 1)
        {
            LOG(INFO) << "[Adaptive] candidates empty ";
            return false;
        }
        return true;
    }

    const Candidate2D &AdaptiveFilter::GetBestCand()
    {
        float scoreMax = 0.0;
        for (size_t i = 0; i < candidates.size(); i++)
        {
            auto tmp = candidates[i].score;
            if (i == 0 || scoreMax < tmp)
            {
                scoreMax = tmp;
                best_idx_ = i;
            }
        }
        const auto &best = candidates[best_idx_];
        offsetBest = Eigen::Vector3d(best.x, best.y, best.orientation);
        LOG(INFO) << "[best] " << best_idx_ << " " << best.x << best.y << common::RadToDeg(best.orientation) << " score " << scoreMax;
        return best;
    }

    void AdaptiveFilter::ScoreCandidates(const Grid2D &grid, const std::vector<DiscreteScan2D> &discrete_scans)
    {
        for (auto &cand : candidates)
        {
            cand.score = ComputeCandidateScore(static_cast<const ProbabilityGrid &>(grid), discrete_scans[cand.scan_index],
                                               cand.x_index_offset, cand.y_index_offset);
            float range = std::hypot(cand.x, cand.y);
            cand.score *= std::exp(-common::Pow2(range * cost_weight_linear + std::abs(cand.orientation) * cost_weight_angular));
        }
    }

    AdaptiveMatcherAngle2D::AdaptiveMatcherAngle2D(const ScanMatchOptions &options)
        : options_(options)
    {
        filter_ = std::make_shared<AdaptiveFilter>(options);
    }

    double AdaptiveMatcherAngle2D::Match(const transform::Rigid2d &poseInit, const sensor::PointCloud &point_cloud, const Grid2D &grid,
                                         transform::Rigid2d *pose_estimate)
    {
        CHECK(pose_estimate != nullptr);
        PoseApi2d pose(poseInit);
        const auto rotated_cloud = sensor::TransformPointCloud(point_cloud, pose.Rotation());

        // 更新参数
        float gridResolution = grid.limits().resolution();
        float rangeMax = ComputeRangeMax(rotated_cloud, gridResolution);
        filter_->UpdateParameter(rangeMax, gridResolution);

        // 在x和y的维采样若干候选位姿，偏移量是指栅格偏移数
        filter_->SampleNears();

        // if (!filter_->Check())
        // {
        //     filter_->inited_ = false;
        //     return 0.0;
        // }

        // 以初始值为基础，生成一些列的旋转后的点云
        const auto rotated_scans = GenerateRotatedScans(rotated_cloud, filter_->rotations_);
        // 根据旋转点云坐标，生成点对应的栅格坐标
        const auto discrete_scans = DiscretizeScans(grid.limits(), rotated_scans, pose.Translation());
        // 计算每个候选位姿的概率得分
        filter_->ScoreCandidates(grid, discrete_scans);
        const auto &bestCand = filter_->GetBestCand();
        // 选出最佳得分的候选位姿作为返回结果
        *pose_estimate = pose.GetPose(bestCand);
        return bestCand.score;
    }

    void AdaptiveFilter::Failed()
    {
        gauss_.Clear();
        inited_ = false;
    }

    void AdaptiveMatcherAngle2D::Failed()
    {
        filter_->Failed();
    }

    int AdaptiveMatcherAngle2D::Success(const transform::Rigid2d &pose_init, const transform::Rigid2d &pose_opt)
    {
        int code = filter_->Success(pose_init, pose_opt);
        if (code > 0)
            filter_->inited_ = true;
        else
            filter_->Failed();
        LOG(INFO) << "[Success] " << code;
        return code;
    }
} // namespace cartographer
